{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ShHGnzOb7Tnr"
      },
      "source": [
        "##### Copyright 2025 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tod4A_Ic7Y3I"
      },
      "outputs": [],
      "source": [
        "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RWf7Xw5BdPmh"
      },
      "source": [
        "# Using Gemini API with Qdrant vector search for hybrid retrieval in legal AI\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VFNQtwrU7iUV"
      },
      "source": [
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemini/cookbook/blob/main/examples/qdrant/Hybrid_Search_Legal.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" height=30/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ntUiUHdJ7kh4"
      },
      "source": [
        "<!-- Community Contributor Badge -->\n",
        "<table>\n",
        "  <tr>\n",
        "    <!-- Author Avatar Cell -->\n",
        "    <td bgcolor=\"#d7e6ff\">\n",
        "      <a href=\"https://github.com/mrscoopers\" target=\"_blank\" title=\"View Jenny's profile on GitHub\">\n",
        "        <img src=\"https://github.com/mrscoopers.png?size=100\"\n",
        "             alt=\"Jenny's GitHub avatar\"\n",
        "             width=\"100\"\n",
        "             height=\"100\">\n",
        "      </a>\n",
        "    </td>\n",
        "    <!-- Text Content Cell -->\n",
        "    <td bgcolor=\"#d7e6ff\">\n",
        "      <h2><font color='black'>This notebook was contributed by <a href=\"https://github.com/mrscoopers\" target=\"_blank\"><font color='#217bfe'><strong>Jenny</strong></font></a>.</font></h2>\n",
        "      <h5><font color='black'><a href=\"https://www.linkedin.com/in/evgeniya-sukhodolskaya/\"><font color=\"#078efb\">Jenny's LinkedIn</font></a></h5></font><br>\n",
        "      <!-- Footer -->\n",
        "      <font color='black'><small><em>Have a cool Gemini example? Feel free to <a href=\"https://github.com/google-gemini/cookbook/blob/main/CONTRIBUTING.md\" target=\"_blank\"><font color=\"#078efb\">share it too</font></a>!</em></small></font>\n",
        "    </td>\n",
        "  </tr>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5CrYV5o27nkx"
      },
      "source": [
        "<!-- Princing warning Badge -->\n",
        "<table>\n",
        "  <tr>\n",
        "    <!-- Emoji -->\n",
        "    <td bgcolor=\"#f5949e\">\n",
        "      <font size=30>⚠️</font>\n",
        "    </td>\n",
        "    <!-- Text Content Cell -->\n",
        "    <td bgcolor=\"#f5949e\">\n",
        "      <h3><font color=black>This notebook requires paid tier rate limits to run properly.<br>  \n",
        "(cf. <a href=\"https://ai.google.dev/pricing#veo2\"><font color='#217bfe'>pricing</font></a> for more details).</font></h3>\n",
        "    </td>\n",
        "  </tr>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3UKx6xco7qbI"
      },
      "source": [
        "## Overview"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pAUtYeKqxeid"
      },
      "source": [
        "![Austronaut judge](https://storage.googleapis.com/qdrant-examples/astronaut_judge_2.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vH-hlYwijbJN"
      },
      "source": [
        "In the legal domain, **accuracy** and **factual correctness** are immensely critical.\n",
        "\n",
        "A Legal AI startup that collaborated with [Qdrant](https://qdrant.tech/) has outlined **an approach to securing both in Legal AI applications** (for example, Retrieval Augmented Generation (RAG)-based or agentic):\n",
        "\n",
        "> *“Turn everything into a retrieval problem where you're retrieving ground truth. If you frame it that way, you don't have to worry about hallucinations, as everything given to the user is grounded in some part of a valid document.”*\n",
        "\n",
        "Truly, many Legal AI businesses require **high-quality retrieval** in their applications. To get there, you need:\n",
        "- The knowledge of the right tools and techniques that increase search relevance;\n",
        "- A well-suited embedding model;\n",
        "- Being ready to experiment!:)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4bba7f79e718"
      },
      "source": [
        "### This notebook\n",
        "\n",
        "In this notebook, you’ll learn how to combine `gemini-embedding-001` with the tools provided by the Qdrant vector search engine to build a **legal QA retrieval pipeline**.\n",
        "\n",
        "You'll learn how to:\n",
        "- Set up a hybrid search (dense + keyword) in Qdrant;\n",
        "- Use [Matryoshka Representations](https://huggingface.co/blog/matryoshka) of Gemini embeddings to trade off quality vs. cost."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cf7042fdfe06"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JzLZlC7h3mdh"
      },
      "source": [
        "### Install SDK\n",
        "\n",
        "- `google-genai` for `gemini-embedding-001` embeddings;\n",
        "- `qdrant-client[fastembed]` - the Qdrant's python client;\n",
        "- HuggingFace `datasets` - to load open sourced legal Q&A datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "JjDBA3EoyMOq"
      },
      "outputs": [],
      "source": [
        "%pip install -q -U \"google-genai>=1.0.0\" qdrant-client[fastembed] datasets"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WN1wCwGZ4NAP"
      },
      "source": [
        "### Set up your API keys:\n",
        "\n",
        "- `GOOGLE_API_KEY`, required for using `gemini-embedding-001` embeddings  \n",
        "  (look up how to generate it [here](https://ai.google.dev/gemini-api/docs/api-key))\n",
        "\n",
        "- `QDRANT_API_KEY` and `QDRANT_URL` from a **free-forever** Qdrant Cloud cluster  \n",
        "(you'll be guided on how to get both in the [Qdrant Cloud UI](https://cloud.qdrant.io/))\n",
        "\n",
        "To run the following cell, your API keys must be stored in a Colab Secret tab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "YiE8lhw913pS"
      },
      "outputs": [],
      "source": [
        "from google.colab import userdata\n",
        "\n",
        "GOOGLE_API_KEY = userdata.get('GOOGLE_API_KEY')\n",
        "QDRANT_API_KEY = userdata.get('QDRANT_API_KEY')\n",
        "QDRANT_URL = userdata.get('QDRANT_URL')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IgR_gMjbr08n"
      },
      "source": [
        "## Step 1: Download the Dataset\n",
        "\n",
        "You'll use one of the [Hugging Face datasets from Isaacus](https://huggingface.co/isaacus), a legal artificial intelligence research company.\n",
        "\n",
        "A common use case in legal AI is a Retrieval-Augmented Generation (RAG) chatbot. To evaluate retrieval performance for such applications, you need a Question-Answer (QA) dataset."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "faf83189542a"
      },
      "source": [
        "### Choosing a Dataset\n",
        "\n",
        "- [Open Australian Legal QA](https://huggingface.co/datasets/isaacus/open-australian-legal-qa) looks interesting. However, all its LLM-generated questions mention the exact name of the legal case, which also appears in the answer. The dataset maps each question to one answer (1:1), making it trivial to build a perfect retriever => not even close to real-life scenarios:)\n",
        "\n",
        "- Instead, let's consider [LegalQAEval](https://huggingface.co/datasets/isaacus/LegalQAEval). It looks more like the kind of questions a user might ask a RAG-based legal chatbot. For example:  \n",
        "  * \"*How are pharmacists regulated in most jurisdictions?*\"\n",
        "  * \"*what is ncts*\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9ad609010e48"
      },
      "source": [
        "#### LegalQAEval\n",
        "\n",
        "This dataset contains ~2400 QA pairs and includes:\n",
        "\n",
        "- `id`: a unique string identifier;\n",
        "- `question`: a natural language question;\n",
        "- `text`: a chunk of text that *may* contain the answer;\n",
        "- `answers`: a list of answers (and their positions within the text), or `null` if the `text` does not have the answer."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6iE0KTZaritk"
      },
      "source": [
        "Load the legal QA corpus; you'll use all available splits."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "R78lyu6RuUp-"
      },
      "outputs": [],
      "source": [
        "from datasets import load_dataset, concatenate_datasets\n",
        "\n",
        "corpus = concatenate_datasets(load_dataset('isaacus/LegalQAEval', split=['val', 'test']))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5IJgg0S5tlBt"
      },
      "source": [
        "### Text chunks deduplication\n",
        "\n",
        "Since the dataset can contain `text` chunks with multiple questions related to them, initially deduplicate `text` fields to not store identical information several times."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "53zW9Tr_mBv1"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "import datasets\n",
        "\n",
        "# Convert the Hugging Face dataset to a pandas DataFrame\n",
        "df = corpus.to_pandas()\n",
        "\n",
        "# Group by 'text' and aggregate 'id' into a list\n",
        "grouped_corpus = df.groupby('text')['id'].apply(list).reset_index().rename(columns={'id': 'ids'})\n",
        "\n",
        "corpus_deduplicated = datasets.Dataset.from_pandas(grouped_corpus)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qmQTufefz7Uj"
      },
      "source": [
        "## Step 2: Define the use case configuration\n",
        "\n",
        "In a typical legal chatbot scenario, users ask a question, and an LLM generates an answer based on a relevant text chunk.\n",
        "\n",
        "To imitate it, you'll need to store in Qdrant numerical representations (embeddings) of `text` chunks.  \n",
        "During retrieval, a `question` will be converted into a numerical representation in the same embedding space. Then, (approximately) the nearest `text` chunk will be found in the vector index.\n",
        "\n",
        "> The Gemini embedding model [supports](https://ai.google.dev/gemini-api/docs/embeddings#supported-task-types) RAG-style Q&A retrieval (task type `QUESTION_ANSWERING`).\n",
        "\n",
        "Now, to fully define our storage configuration, let's consider several factors relevant to a common RAG use case in the legal AI domain."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9a00c136f2af"
      },
      "source": [
        "### Cost versus accuracy: matryoshka representations\n",
        "\n",
        "Gemini `gemini-embedding-001` embeddings are 3072-dimensional.  \n",
        "In a RAG setup with ~1 million chunks, storing such embeddings in RAM (for fast retrieval) would require about **12 GB**.\n",
        "\n",
        "The Gemini embedding model supports an approach to balance accuracy & cost of retrieval. It is trained using [Matryoshka Representation Learning (MRL)](https://ai.google.dev/gemini-api/docs/embeddings#control-embedding-size), meaning that the most important information about the encoded text is stored in the first dimensions of the embedding.\n",
        "\n",
        "So, you can, for example:\n",
        "- Use only the first 768 dimensions of the Gemini embedding for **faster retrieval**;\n",
        "- And then **rerank** the retrieved results using the full 3072-dimensional embeddings for higher precision."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f112792b8148"
      },
      "source": [
        "### Accuracy from the best of both worlds: hybrid search\n",
        "\n",
        "In legal use cases, it is often beneficial to combine the strengths of:  \n",
        "- **Keyword-based search (lexical)** for more direct control over matches;  \n",
        "- **Embedding-based search (semantic)** for handling questions phrased in a conversational way.\n",
        "\n",
        "> **In Qdrant, both approaches can be combined in [hybrid & multi-stage queries](https://qdrant.tech/documentation/concepts/hybrid-queries/)**.  \n",
        "\n",
        "For the keyword-based part, Qdrant supports multiple options, from traditional BM25 to sparse neural retrievers like [SPLADE](https://qdrant.tech/documentation/fastembed/fastembed-splade/). Among the options, there's [**our custom improvement of BM25 called miniCOIL**](https://qdrant.tech/documentation/fastembed/fastembed-minicoil/), which you will use in this notebook.\n",
        "\n",
        "> In Qdrant, keyword-based retrieval is achieved using [sparse vectors](https://qdrant.tech/documentation/concepts/vectors/#sparse-vectors).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U1_0T7gHx8Et"
      },
      "source": [
        "### Collection configuration\n",
        "Configure a Qdrant collection for the legal QA retrieval pipeline."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "JgiAtOUMx7HV"
      },
      "outputs": [],
      "source": [
        "from qdrant_client import QdrantClient, models\n",
        "\n",
        "qdrant_client = QdrantClient(  # Initializing Qdrant client.\n",
        "    url=QDRANT_URL,\n",
        "    api_key=QDRANT_API_KEY,\n",
        ")\n",
        "\n",
        "COLLECTION_NAME = \"legal_AI_QA\"\n",
        "GEMINI_EMBEDDING_RETRIEVAL_SIZE = 768\n",
        "GEMINI_EMBEDDING_FULL_SIZE = 3072\n",
        "\n",
        "if not qdrant_client.collection_exists(collection_name=COLLECTION_NAME):\n",
        "    qdrant_client.create_collection(\n",
        "        collection_name=COLLECTION_NAME,\n",
        "        vectors_config={\n",
        "            \"gemini_embedding_retrieve\": models.VectorParams(\n",
        "                size=GEMINI_EMBEDDING_RETRIEVAL_SIZE,  # Smaller embeddings for faster retrieval.\n",
        "                distance=models.Distance.COSINE,\n",
        "            ),\n",
        "            \"gemini_embedding_rerank\": models.VectorParams(\n",
        "                size=GEMINI_EMBEDDING_FULL_SIZE,  # Full-sized embeddings for precision-boosting reranking.\n",
        "                distance=models.Distance.COSINE,\n",
        "                hnsw_config=models.HnswConfigDiff(\n",
        "                    m=0  # Since these embeddings aren't used for retrieval, you don't need to spend resources on building a vector index.\n",
        "                ),\n",
        "                on_disk=True,  # To save on RAM used for retrieval.\n",
        "            ),\n",
        "        },\n",
        "        sparse_vectors_config={\n",
        "            \"miniCOIL\": models.SparseVectorParams(\n",
        "                modifier=models.Modifier.IDF  # Inverse Document Frequency statistic, computed on the Qdrant side.\n",
        "            )\n",
        "        },\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wbpb6r53K-U5"
      },
      "source": [
        "## Step 3: embed texts & index data to Qdrant\n",
        "\n",
        "To speed up the process of converting the data, you'll:\n",
        "\n",
        "1. Embed with Gemini all `text` chunks in batches using the `get_embeddings_batch` function.  \n",
        "2. Upload the results to Qdrant in batches.  \n",
        "The Qdrant Python client provides the functions `upload_collection` and `upload_points`. These handle batching, retries, and parallelization. They take generators as input, so you'll create a generator function `qdrant_points_stream` for this purpose.\n",
        "\n",
        "> **Note:** Qdrant automatically normalizes uploaded embeddings if the distance function in your collection was set to `COSINE` (cosine similarity). This means you don’t need to pre-normalize truncated Gemini Matryoshka embeddings, as it's [recommended in the Gemini documentation](https://ai.google.dev/gemini-api/docs/embeddings#quality-for-smaller-dimensions).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5883cd1948e2"
      },
      "outputs": [],
      "source": [
        "GEMINI_MODEL_ID = \"gemini-embedding-001\" # @param [\"gemini-embedding-001\"] {\"allow-input\":true, isTemplate: true}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "DJeUdy7OybeX"
      },
      "outputs": [],
      "source": [
        "from google import genai\n",
        "from google.genai import types\n",
        "from google.api_core import retry\n",
        "import uuid\n",
        "\n",
        "google_client = genai.Client(api_key=GOOGLE_API_KEY)\n",
        "\n",
        "@retry.Retry(timeout=300)\n",
        "def get_embeddings_batch(texts, task_type: str = \"RETRIEVAL_DOCUMENT\"):\n",
        "    \"\"\"Generates embeddings for a batch of texts.\n",
        "\n",
        "    Args:\n",
        "        texts: A list of strings to embed.\n",
        "        task_type: The task type for the embedding model.\n",
        "\n",
        "    Returns:\n",
        "        A list of embedding vectors.\n",
        "\n",
        "    Raises:\n",
        "        Exception: If an error occurs during embedding generation.\n",
        "    \"\"\"\n",
        "    try:\n",
        "        res = google_client.models.embed_content(\n",
        "            model=GEMINI_MODEL_ID,\n",
        "            contents=texts,\n",
        "            config=types.EmbedContentConfig(task_type=task_type),\n",
        "        )\n",
        "        return [e.values for e in res.embeddings]\n",
        "    except Exception as e:\n",
        "        print(f\"An error occurred while getting embeddings: {e}\")\n",
        "        raise\n",
        "\n",
        "\n",
        "def qdrant_points_stream(corpus, avg_corpus_text_length, gemini_batch_size: int = 8):\n",
        "    \"\"\"Streams Qdrant points with embeddings for a given corpus.\n",
        "\n",
        "    Args:\n",
        "        corpus: The dataset to process.\n",
        "        avg_corpus_text_length: The average text length for miniCOIL (based on BM25 formula).\n",
        "        gemini_batch_size: The batch size for Gemini embedding requests.\n",
        "\n",
        "    Yields:\n",
        "        Qdrant PointStruct objects.\n",
        "    \"\"\"\n",
        "    for start in range(0, len(corpus), gemini_batch_size):  # Iterate over the dataset in batches.\n",
        "        end = min(start + gemini_batch_size, len(corpus))\n",
        "        batch = corpus.select(range(start, end))  # Current batch slice.\n",
        "\n",
        "        gemini_embeddings_full = get_embeddings_batch(\n",
        "            [row[\"text\"] for row in batch], task_type=\"RETRIEVAL_DOCUMENT\"\n",
        "            )  # Generate embeddings for this batch.\n",
        "\n",
        "        for batch_item, gemini_embedding_full in zip(batch, gemini_embeddings_full):\n",
        "            yield models.PointStruct(\n",
        "                id=str(uuid.uuid4()),  # Unique ID (string UUID or integer supported by Qdrant).\n",
        "                payload={  # Metadata stored alongside the vector.\n",
        "                    \"text\": batch_item[\"text\"],  # Raw text for users/LLMs.\n",
        "                    \"ids\": batch_item[\"ids\"],  # IDs of the QA pairs related to this `text` (for later evaluation).\n",
        "                },\n",
        "                vector={  # Embeddings.\n",
        "                    \"gemini_embedding_rerank\": gemini_embedding_full,  # Full Gemini embedding for reranking.\n",
        "                    \"gemini_embedding_retrieve\": gemini_embedding_full[:768],  # Truncated Gemini embedding for retrieval.\n",
        "                    \"miniCOIL\": models.Document(  # Custom Qdrant-optimized BM25 replacement.\n",
        "                        text=batch_item[\"text\"],\n",
        "                        model=\"Qdrant/minicoil-v1\",\n",
        "                        options={\"avg_len\": avg_corpus_text_length, \"k\": 0.9, \"b\": 0.4},  # Corpus avg length, k_1 & b from BM25 formula.\n",
        "                    ),\n",
        "                },\n",
        "            )\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ECrXnkXFNnui"
      },
      "source": [
        "Now you'll embed the data and upload the embeddings.\n",
        "\n",
        "> Try experimenting with different batch sizes when generating embeddings and uploading them to Qdrant.  \n",
        "The fastest setup usually depends on your network speed & RAM/CPU/GPU, and keep in mind that embedding inference is not a very fast process.\n",
        "\n",
        "> The representations used in Qdrant for the keyword-based retrieval part of hybrid search are produced by Qdrant.  \n",
        "In Colab, Qdrant will download the required models the first time you use them (in our case, **Qdrant/minicoil-v1**), as they’re needed for converting `text` chunks to sparse representations.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "uxferMA1p1vn"
      },
      "outputs": [],
      "source": [
        "import tqdm\n",
        "\n",
        "COLLECTION_NAME = \"legal_AI_QA\"\n",
        "\n",
        "# Estimating the average length of the texts in the corpus on the subsample of 1000, to use in BM25-inspired keywords-based retrieval.\n",
        "SUBSET_SIZE = 1000\n",
        "avg_corpus_text_length = sum(len(text.split()) for text in corpus[\"text\"][:SUBSET_SIZE]) / SUBSET_SIZE\n",
        "\n",
        "qdrant_client.upload_points(\n",
        "    collection_name=COLLECTION_NAME,\n",
        "    points=tqdm.tqdm(\n",
        "        qdrant_points_stream(corpus_deduplicated,\n",
        "                            avg_corpus_text_length=avg_corpus_text_length,\n",
        "                            gemini_batch_size=4),\n",
        "        desc=\"Uploading points\",\n",
        "    ),\n",
        "    batch_size=4,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2vGGgWv2_nvK"
      },
      "source": [
        "## Step 4: experiment & evaluate\n",
        "\n",
        "What’s important for every retrieval task is experimenting with different instruments & running evaluations based on a sensible metric."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0d10b3e83010"
      },
      "source": [
        "### Metric\n",
        "In RAG, the goal is usually to get the **correct result within the top-N retrieved results, using a very small N**, since that’s what the LLM will use to generate a grounded answer, and you'd want to save context window size/reduce token costs.  \n",
        "\n",
        "You'll use the metric **`hit@1`**, meaning the top-1 ranked text chunk is actually the answer to the question."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9d63800b156e"
      },
      "source": [
        "### Eval set\n",
        "For experiments, you should only use questions where the `answers` field is not `null`, since this guarantees that this text chunk contains the answer to the question.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "KnuRkPJkGv_c"
      },
      "outputs": [],
      "source": [
        "questions = corpus.filter(lambda item: len(item['answers']) > 0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GDvQtFYkS-c1"
      },
      "source": [
        "Inference Gemini embeddings for all the questions, so you can experiment freely without spending extra time or money."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "JlvBnckV1UYh"
      },
      "outputs": [],
      "source": [
        "import tqdm\n",
        "\n",
        "question_embeddings = {}\n",
        "\n",
        "question_texts = [q['question'] for q in questions]\n",
        "question_ids = [q['id'] for q in questions]\n",
        "all_embeddings = []\n",
        "\n",
        "BATCH_SIZE = 32\n",
        "\n",
        "for i in tqdm.tqdm(range(0, len(question_texts), BATCH_SIZE), desc=\"Embedding questions\"):\n",
        "    batch_texts = question_texts[i:i + BATCH_SIZE]\n",
        "    embeddings = get_embeddings_batch(batch_texts, task_type=\"QUESTION_ANSWERING\")\n",
        "    all_embeddings.extend(embeddings)\n",
        "\n",
        "question_embeddings = {qid: emb for qid, emb in zip(question_ids, all_embeddings)}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m6UoRhLATGlZ"
      },
      "source": [
        "And randomly select a test subset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "tl46lbzd8nv0"
      },
      "outputs": [],
      "source": [
        "questions = questions.shuffle(seed=42).select(range(500))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QEG_H0UlI1L9"
      },
      "source": [
        "### Experiment\n",
        "\n",
        "There are many ways to improve search results. For example, **reranking** alone can be done with high-dimensional embeddings like Gemini, [multivectors like ColBERT](https://qdrant.tech/documentation/advanced-tutorials/using-multivector-representations/) or cross-encoders.\n",
        "\n",
        "For simplicity, you'll focus on three simple retrieval approaches, three experiments that are a good starting point for high-precision-demanding domains like legal:\n",
        "\n",
        "**Experiment 1: Vanilla Retrieval**  \n",
        "Use truncated Gemini embeddings for vanilla retrieval.  \n",
        "This gives you a simple reference point to compare improvements against.  \n",
        "\n",
        "**Experiment 2: Reranking**  \n",
        "Rerank the retrieved subset with full-sized Gemini embeddings.  \n",
        "Larger embeddings capture finer semantic details that the retriever may miss.\n",
        "\n",
        "**Experiment 3: Hybrid Search**  \n",
        "Combine semantic (captures meaning) and keyword-based (ensures exact matches) retrieval in **Hybrid Search**.  \n",
        "*For this toy dataset, keyword matching may not add much, as all questions are very \"conversational\" style with not-so-many-keywords, but in real-life legal AI retrieval, it makes a difference*\n",
        "\n",
        "The setup is the following:\n",
        "1. Run two searches with the same query.  \n",
        "2. Merge the results into a single list with a fusion algorithm. Here you'll use [Reciprocal Rank Fusion (RRF)](https://qdrant.tech/documentation/concepts/hybrid-queries/#hybrid-search), a simple well-known zero-shot method of fusion.  \n",
        "\n",
        "---\n",
        "Compare the results of all three experiments on our evaluation set using the chosen metric.  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_Ke-sa9UWSOD"
      },
      "outputs": [],
      "source": [
        "COLLECTION_NAME = \"legal_AI_QA\"\n",
        "N = 1\n",
        "\n",
        "def hit_at_n(results: list[models.ScoredPoint], question_id: str, N: int) -> int:\n",
        "    \"\"\"Calculates if the correct document is within the top N retrieved results.\n",
        "\n",
        "    Args:\n",
        "        results: A list of scored points from a Qdrant search.\n",
        "        question_id: The ID of the question to check for.\n",
        "        N: The number of top results to check.\n",
        "\n",
        "    Returns:\n",
        "        1 if the correct document is found, 0 otherwise.\n",
        "    \"\"\"\n",
        "    for result in results:\n",
        "        if question_id in result.payload[\"ids\"]:\n",
        "            return 1\n",
        "    return 0\n",
        "\n",
        "hits_baseline = 0\n",
        "hits_rerank = 0\n",
        "hits_hybrid = 0\n",
        "\n",
        "for question in tqdm.tqdm(questions, desc=f\"Evaluating hits@{N}\"):\n",
        "    full_gemini_embedding = question_embeddings[question['id']] # Embedding of the question.\n",
        "\n",
        "    ## Experiment 1: Baseline\n",
        "    result_baseline = qdrant_client.query_points(\n",
        "        collection_name=COLLECTION_NAME,\n",
        "        query=full_gemini_embedding[:768],  # Use the first quarter of the Gemini embedding for retrieval.\n",
        "        limit=N,\n",
        "        using=\"gemini_embedding_retrieve\",\n",
        "        with_payload=True,\n",
        "    )\n",
        "    hits_baseline += hit_at_n(result_baseline.points, question_id=question['id'], N=N)\n",
        "\n",
        "    ## Experiment 2: Reranking\n",
        "    result_rerank = qdrant_client.query_points(\n",
        "        collection_name=COLLECTION_NAME,\n",
        "        prefetch=models.Prefetch( # First, retrieve 50 candidates with the smaller embedding.\n",
        "            query=full_gemini_embedding[:768],\n",
        "            using=\"gemini_embedding_retrieve\",\n",
        "            limit=50\n",
        "        ),\n",
        "        query=full_gemini_embedding, # Then rerank those 50 results with the full embedding.\n",
        "        limit=N,\n",
        "        using=\"gemini_embedding_rerank\",\n",
        "        with_payload=True\n",
        "    )\n",
        "    hits_rerank += hit_at_n(result_rerank.points, question_id=question['id'], N=N)\n",
        "\n",
        "    ## Experiment 3: Hybrid search (semantic + keyword)\n",
        "    result_hybrid = qdrant_client.query_points(\n",
        "        collection_name=COLLECTION_NAME,\n",
        "        prefetch=[\n",
        "            models.Prefetch( # Retrieve 25 results using semantic search (Gemini truncated embeddings).\n",
        "                query=full_gemini_embedding[:768],\n",
        "                using=\"gemini_embedding_retrieve\",\n",
        "                limit=25\n",
        "            ),\n",
        "            models.Prefetch( # Retrieve 25 results using miniCOIL (Qdrant’s custom improved version of BM25-based keyword search).\n",
        "                query=models.Document(\n",
        "                    text=question['question'],\n",
        "                    model=\"Qdrant/minicoil-v1\"\n",
        "                ),\n",
        "                using=\"miniCOIL\",\n",
        "                limit=25\n",
        "            )\n",
        "        ],\n",
        "        query=models.FusionQuery(fusion=\"rrf\"), # Fuse the two result sets with Reciprocal Rank Fusion (RRF).\n",
        "        limit=N,\n",
        "        with_payload=True\n",
        "    )\n",
        "    hits_hybrid += hit_at_n(result_hybrid.points, question_id=question['id'], N=N)\n",
        "\n",
        "\n",
        "hits_at_N_baseline = hits_baseline / len(questions) # Compute average hits@N for each experiment.\n",
        "hits_at_N_rerank = hits_rerank / len(questions)\n",
        "hits_at_N_hybrid = hits_hybrid / len(questions)\n",
        "\n",
        "print(\"\\n\")\n",
        "print(f\"Retrieval Avg Hits@{N}: {hits_at_N_baseline:.3f} ({hits_baseline}/{len(questions)})\")\n",
        "print(f\"Rerank Avg Hits@{N}: {hits_at_N_rerank:.3f} ({hits_rerank}/{len(questions)})\")\n",
        "print(f\"Hybrid Search Avg Hits@{N}: {hits_at_N_hybrid:.3f} ({hits_hybrid}/{len(questions)})\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yyEngxhDUZi_"
      },
      "source": [
        "## Next steps\n",
        "\n",
        "In this notebook, you set up a **retrieval pipeline behind a typical legal RAG chatbot** with **Qdrant Vector Search Engine** and **Gemini Embeddings**.   \n",
        "You tried several approaches to retrieval, making use of Gemini's capability to generate Matryoshka representations & Qdrant's tooling for retrieval with reranking and hybrid search.\n",
        "\n",
        "Of course, legal applications require much more than a plain zero-shot pipeline. Retrieval quality always depends on the dataset and use case, so there’s no silver bullet besides experimenting & iterating."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "575d76e0b4df"
      },
      "source": [
        "### Where to go from zero-shot:\n",
        "- Analyze queries-misses;  \n",
        "- Tune vector index parameters (for example, `ef` for search at scale in Qdrant);  \n",
        "- Experiment with different fusion strategies & parameters in hybrid search;  \n",
        "- Try query expansion (or filters extraction).  \n",
        "- ...\n",
        "\n",
        "Use this notebook as a baseline to build on and experiment to find what works best for you!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DGZUuhJyD3tC"
      },
      "source": [
        "![Austronaut judge 2](https://storage.googleapis.com/qdrant-examples/astronaut_court_1.png)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "Hybrid_Search_Legal.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
